# -*- coding: utf-8 -*-
"""
Created on Wed Dec 16 20:33:39 2020

@author: Team317
"""

#   模型预加载数据的预处理部分
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms, models
from torch.utils.data import DataLoader
from age_net import AgeNet
import matplotlib.pyplot as plt
import torch.nn.functional as F
from time import time
import os
import numpy as np
from torch.utils.data.sampler import WeightedRandomSampler
os.environ['CUDA_VISIBLE_DEVICE']='1'
class AgeTrain:
    
    def __init__(self, model_path = '../model/MSE_model_001.pth'):
        self.device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')  # 处理器选择
        agenet = AgeNet.agenet(path)
        self.agenet = agenet.to(self.device)
        
        
        # 参数全部参与训练
        for param in self.agenet.parameters():
            param.requires_grad = True
        
        # 参数设置
        LR = 1e-3  # 学习率
        print('device:',self.device)
        # 平方差损失函数，计算二范数
        self.MSE_loss = nn.MSELoss()
        # 定义优化器
        self.optimizer = optim.Adam(self.agenet.parameters(), LR)
        
        
    
    def train(self, epoch, save_path, model_path, batch=1000 ):
        if os.path.exists(model_path):
            print("存在已训练的模型，正在加载{}...".format(model_path))
            self.agenet.load_state_dict(torch.load(model_path))
        transform = transforms.Compose([
            transforms.Resize(112),
            transforms.ToTensor()
            # transforms.ColorJitter(0.2,0.2,0.2,0.01),
            # transforms.Normalize((0.5, 0.5, 0.5), (128 / 255., 128 / 255., 128 / 255.))
        ])
        # 使用WeightRandomSampler来平衡数据集
        samples_count = [1647,2845,1512,4900,4900,4900,4900,2866,1382,412]
        samples_weight = 1/(samples_count/np.sum(samples_count))
        samples_set = [samples_weight[k] for k in range(0,len(samples_count)) for i in range(0,samples_count[k])]
        samples_num = 30264
        sampler = WeightedRandomSampler(samples_set,samples_num)
        
        root = '../trainSet'
        train_dataset = datasets.ImageFolder(root, transform)
        # train_loader = torch.utils.data.DataLoader(train_dataset,num_workers=8, batch_size=8, shuffle=True)
        train_loader = torch.utils.data.DataLoader(train_dataset,num_workers=8, batch_size=8, sampler=sampler)

        self.agenet.train()
        loss_list = []
        bias_sum = 0
        loss_sum = 0
        for i, data in enumerate(train_loader):
            
            # data = getImage(root_path = '../AFAD_image')
            
            # 获得数据和对应的标签
            inputs, labels = data
            inputs = inputs.to(self.device)
            labels = labels.to(self.device)
            inputs.requires_grad = True
            
            # 获得模型预测结果
            out = self.agenet(inputs).view(labels.size())
    
            # 计算二范数loss
            loss = self.MSE_loss(out.float(), labels.float())
            loss_list.append(loss.item())
            
            # 计算梯度
            loss.backward()
            # 修改权值
            self.optimizer.step()
            # 梯度清0
            self.optimizer.zero_grad()
            
            
            # 显示偏差和loss
            bias_sum += abs(labels-out).sum().item()
            loss_sum += loss.item()
            times = 50
            if (i+1) % times == 0:
                print("第{}次训练".format(i), end='\t')
                # 显示loss和标签
                loss_aver = loss_sum / times
                labels_data = [i.item() for i in labels]
                pre_data = ['{0:.2f}'.format(i.item()) for i in out]
                print('loss_aver:{0:.4f} \n  label:{1}\n  pre:{2}'.format(loss_aver, labels_data, pre_data))
                
                # 显示偏差
                # bias = self.MSE_bias(train_loader)
                bias_aver = bias_sum/times/8
                print('bias_aver:', bias_aver)
                bias_sum = 0  # 偏差值归零
                loss_sum = 0
                
                
            
            if i>=batch:
                break
        # 展示loss变化 
        self.showChart(loss_list,epoch)
        torch.save(self.agenet.state_dict(), save_path)
        return loss_list
        
    
    def test(self, model_path):
        try:
            self.agenet.load_state_dict(torch.load(model_path))
        except:
            print("模型不存在!")

        root = '../testSet'
        transform = transforms.Compose([
            transforms.Resize(112),
            transforms.ToTensor()
        ])
        test_dataset = datasets.ImageFolder(root, transform)
        test_loader = torch.utils.data.DataLoader(test_dataset,num_workers=8, batch_size=8, shuffle=True)

        self.agenet.eval()
        loss_list = []
        bias_sum = 0
        loss_sum = 0
        count = 0
        for i, data in enumerate(test_loader):
            
            # data = getImage(root_path = '../AFAD_image')
            
            # 获得数据和对应的标签
            inputs, labels = data
            inputs = inputs.to(self.device)
            labels = labels.to(self.device)
            inputs.requires_grad = False
            
            # 获得模型预测结果
            out = self.agenet(inputs).view(labels.size())     
            
            
            # 显示偏差和loss
            bias = abs(labels-out).sum().item()/8
            bias_sum += bias*8
            
            if bias>1:
                count+=1
                print("labels:{}  pre:{}".format(labels, out))
            
        # 显示偏差
        bias_aver = bias_sum/8/(i+1)
        print('bias_aver in test set:', bias_aver)
        print("count",count)

    
    
    def showChart(self, data_list,ep):
        x = [i for i in range(len(data_list))]
        plt.title(ep)
        plt.plot(x, data_list)
        plt.show()

if __name__ == '__main__':
    print("模型加载...")
    path = '../Arcface_100.pth'
    
    loss_list = []  # 记录损失值的变化
    acc_list = []  # 记录准去率的变化
    start_epoch = 100
    end_epoch = 120
    ageTrain = AgeTrain()
    # # 读取数据
    for ep in range(start_epoch, end_epoch):
        print("第{}个epoch:".format(ep))
        save_path = '../model/tanh_{0:3d}'.format(ep)
        model_path = '../model/tanh_{0:3d}'.format(ep-1)
        
        ageTrain.train(ep, save_path, model_path, batch=2000)
    
    # model_path = '../model/tanh_{0:3d}'.format(start_epoch-1)
    # ageTrain.test(model_path)
        
    
        

        
